'''
PdfProcessor类
用于提取PDF文件中的文本、表格、图片等信息
'''
import os
import re
import fitz
import time
import random
import string
import datetime
import traceback
import pdfplumber
import numpy as np
from common.custom.ocr import MyOCR
from multiprocessing import Pool
from pathos.multiprocessing import Pool as PathosPool
# import guppy


class PdfProcessor():
    '''
    描述：
        PdfProcessor类，利用PaddleOCR用于提取PDF文件中的文本、图像、表格
    参数：
        filepath: PDF文件路径
        media_root: 保存图片的路径
    成员变量：
        document_info: list(page_info)
        page_info: {
            "pno": 页码,
            "type": 页面类型,
            "content": 页面内容
        }
    '''
    def __init__(self, filepath, media_root) -> None:
        self.filepath = filepath # PDF文件路径
        try:
            self.documnet = fitz.open(filepath) # PyMuPdf打开PDF文件
            self.pdfplumber = pdfplumber.open(filepath) # pdfplumber打开PDF文件
            self.media_root = media_root # 保存图片的路径
            self.pdf_ocr = MyOCR()
            self.zoom_x = 2.0 # 缩放比例
            self.zoom_y = 2.0 # 缩放比例
            self.mat = fitz.Matrix(self.zoom_x, self.zoom_y) # 缩放矩阵
            self.y_threshold = 300 # y轴阈值
            self.x_threshold = 20 # x轴阈值
            self.processing_number = 4 # 处理的进程数
            self.filename = os.path.basename(filepath) # 文件名
            # self.h = guppy.hpy()
            self.message = ""
        except Exception as e:
            print(f"PdfProcessor init error: {e}")
            traceback.print_exc()
            # write_to_log(f"Pdf {filepath} init error: {e}")
            # write_to_log(traceback.format_exc())
            self.message += f"Pdf {filepath} init error: {e}"
            self.message += traceback.format_exc()
            raise Exception(f"Pdf {filepath} init error: {e}")

    def run(self):
        """
        描述: PDF处理
        """
        try:
            self.document_info = [] # PDF每一页的信息
            self.img_save_paths = [] # 保存图片的路径
            self.confidence_page_list = [] # PDF每一页的置信度
            for pno, page in enumerate(self.documnet):
                rect = page.rect
                if rect.width / rect.height <= 1.4:
                    page_info = self.single_page(pno, page) # 单页
                else:
                    page_info = self.double_page(pno, page) # 双页
                self.document_info.append(page_info) # 添加到文档信息中
                self.confidence_page_list.append(page_info["confidence_page"]) # 添加到置信度列表中
            
            # print(self.h.heap())
            # self.delete_images(self.img_save_paths) # 删除临时图片

            self.pdfplumber.close() # 关闭pdfplumber

            # 置信度
            # 去除self.confidence_page_list中nan和非数字的元素
            self.confidence_page_list = [x for x in self.confidence_page_list if (str(x) != 'nan' and x is not None)]

            if len(self.confidence_page_list):
                self.confidence = np.mean(self.confidence_page_list) # PDF整体平均置信度
                self.confidence_max = np.max(self.confidence_page_list) # PDF整体最大置信度
                self.confidence_min = np.min(self.confidence_page_list) # PDF整体最小置信度
            else:
                self.confidence = 0
                self.confidence_max = 0
                self.confidence_min = 0

        except Exception as e:
            print(f"PdfProcessor run error: {e}")
            traceback.print_exc()
            # write_to_log(traceback.format_exc())
            self.message += f"PdfProcessor run error: {e}"
            self.message += traceback.format_exc()
            raise Exception(f"PdfProcessor run error: {e}")

    def run_multiprocessing(self):
        """
        描述: 多进程PDF处理
        # FIXME: 有bug，待修复
        """
        self.img_save_paths = [] # 保存图片的路径
        self.mutil_process_res = [] # 进程池中的结果
        processing_pool = PathosPool(self.processing_number) # 创建进程池
        for pno, page in enumerate(self.documnet):
            rect = page.rect
            if rect.width / rect.height <= 1.4:
                page_info = processing_pool.apply_async(self.single_page, args=(pno, page,))
            else:
                page_info = processing_pool.apply_async(self.double_page, args=(pno, page,))
            self.mutil_process_res.append(page_info)
            time.sleep(0.5)
        
        processing_pool.close() # 关闭进程池
        processing_pool.join() # 等待进程池中的进程执行完毕
        self.document_info = [res.get() for res in self.mutil_process_res] # 获取进程池中的结果
        self.delete_images(self.img_save_paths) # 删除临时图片
        self.pdfplumber.close() # 关闭pdfplumber

    def single_page(self, pno, page):
        '''
        描述：
            单页PDF处理
        参数:
            pno: 页码
            page: PyMuPdf的Page对象
        返回值：
            page_info: 页面信息
        '''
        # 保存PDF页面图片
        now = datetime.datetime.now().strftime('%Y%m%d%H%M%S')
        suffix = ''.join(random.sample(string.ascii_letters + string.digits, 8))
        img_save_path = os.path.join(self.media_root, 'temp_images', f"{self.filename}_images_{pno}_{suffix}_{now}.png")
        page.get_pixmap(matrix=self.mat).save(img_save_path)
        self.img_save_paths.append(img_save_path)

        # 利用 PPStructure 和 paddleocr 进行版面分析和文字提取
        structure = self.pdf_ocr.get_structure(img_save_path) # 速度比较慢
        error_axis_x = 50 if self.is_single_colum(structure=structure) else 5 # 单栏双栏判断

        # 获取页面内容和置信度 by PaddleOCR [速度慢]
        content, confidence_page = self.get_content_confidence_by_PaddleOCR(structure, img_save_path, error_axis_x) 

        page_info = {"pno": pno} # 页面信息
        page_info["content"] = content
        page_info["confidence_page"] = confidence_page
        page_info["image_count"] = self.get_image_count(structure) # 图片数量
        page_info["table_count"] = self.get_table_count(structure) # 表格数量
        page_info["new_structure"] = self.get_image_table_count(structure) # 每一块下方的图片数量和表格数量

        self.delete_image(img_save_path)

        return page_info

    def double_page(self, pno, page):
        '''
        描述：
            双页PDF处理
        参数:
            pno: 页码
            page: PyMuPdf的Page对象
        返回值：
            page_info: 页面信息
        '''
        # PDF页面图片存储路径
        now = datetime.datetime.now().strftime('%Y%m%d%H%M%S')
        suffix = ''.join(random.sample(string.ascii_letters + string.digits, 8))
        img_left_save_path = os.path.join(self.media_root, 'temp_images', f"{self.filename}_images_left_{pno}_{suffix}_{now}.png")
        img_right_save_path = os.path.join(self.media_root, 'temp_images', f"{self.filename}_images_right_{pno}_{suffix}_{now}.png")
        self.img_save_paths.append(img_left_save_path)
        self.img_save_paths.append(img_right_save_path)

        rect = page.rect # 页面矩形
        middle_top_point = fitz.Point((rect.br[0] + rect.tl[0]) * 0.5, rect.tl[1]) # 中线上点
        middle_bottom_point = fitz.Point((rect.br[0] + rect.tl[0]) * 0.5, rect.br[1]) # 中线下点
        clip_left = fitz.Rect(rect.tl, middle_bottom_point) # 左边矩形
        clip_right = fitz.Rect(middle_top_point, rect.br) # 右边矩形
        
        # 保存图片
        page.get_pixmap(matrix=self.mat, clip=clip_left).save(img_left_save_path)
        page.get_pixmap(matrix=self.mat, clip=clip_right).save(img_right_save_path) 

        # 利用 PPStructure 和 paddleocr 进行版面分析和文字提取
        structure_left = self.pdf_ocr.get_structure(img_left_save_path)
        structure_right = self.pdf_ocr.get_structure(img_right_save_path)

        error_axis_x_left = 50 if self.is_single_colum(structure=structure_left) else 5 # 单栏双栏判断
        error_axis_x_right = 50 if self.is_single_colum(structure=structure_right) else 5 # 单栏双栏判断

        # 获取页面内容和置信度 by PaddleOCR [速度慢]
        content_left, confidence_page_left = self.get_content_confidence_by_PaddleOCR(structure_left, img_left_save_path, error_axis_x_left)
        content_right, confidence_page_right = self.get_content_confidence_by_PaddleOCR(structure_right, img_right_save_path, error_axis_x_right)
        content = content_left + content_right
        confidence_page = (confidence_page_left + confidence_page_right) / 2
        image_count = self.get_image_count(structure_left) + self.get_image_count(structure_right)
        table_count = self.get_table_count(structure_left) + self.get_table_count(structure_right)
        new_structure = self.get_image_table_count(structure_left) + self.get_image_table_count(structure_right)

        page_info = {
            "pno": pno,
            "content": content,
            "image_count": image_count,
            "table_count": table_count,
            "new_structure": new_structure,
            "confidence_page": confidence_page
        } 

        self.delete_image(img_left_save_path)
        self.delete_image(img_right_save_path)
        return page_info

    def get_content_by_PPStructure(self, structure):
        """
        描述：
            使用 PPStructure 获取文本内容
        参数：
            structure: 版面分析结果 List[Dict]
        返回值：
            content: 文本内容
        """
        content = ""
        for item in structure:
            if item["type"] == "text":
                for line in item["res"]:
                    content += line["text"]
        content = clean_content(content)
        return content

    def get_content_confidence_by_PaddleOCR(self, structure, img_save_path, error_axis_x):
        """
        描述：
            使用 PaddleOCR + PPStructure 获取文本内容和置信度
            精度高，速度慢
        参数：
            structure: 版面分析结果 List[Dict]
        返回值：
            content: 文本内容
            confidence_page: 置信度
        """
        # 获取paddle ocr的识别结果
        ocr_result = self.pdf_ocr.get_ocr_result(img_save_path)

        # 获取文字部分的Bbox
        text_bboxes = self.get_text_bboxes(structure)

        content = "" # 文本内容
        confidence_list = [] # 置信度
        # 遍历已排序完成的structure的text_bboxes
        for text_bbox in text_bboxes:
            block_res = [] # 存放块中的内容
            for line in ocr_result:
                # line的结构：[[[左上x, 左上y], [右上x, 右上y], [右下x, 右下y], [左下x, 左下y]], (文本, 置信度)]
                # 提取line的bbox [左上x, 左上y, 右下x, 右下y]
                bbox = [line[0][0][0], line[0][0][1], line[0][2][0], line[0][2][1]]
                # 判断line的bbox是否在text_bbox中
                flag = self.is_in_bboxes(bbox, [text_bbox], error_axis_x=error_axis_x, error_axis_y=2)
                if flag:
                    block_res.append(line)
                
                confidence = line[1][1] # 当前line的置信度
                confidence_list.append(confidence)
            
            # 对块内的内容按Y轴排序
            block_res = sorted(block_res, key=lambda x: ((x[0][0][1] + x[0][2][1])*0.5))
            content += "".join([line[1][0] for line in block_res])
    
        content = clean_content(content)
        confidence_page = np.mean(confidence_list) # 当前页面的平均置信度
        return content, confidence_page

    def get_image_count(self, structure):
        """
        描述：获取图片数量
        参数：
            structure: 版面分析结果 List[Dict]
        返回值：
            image_count: 图片数量
        """
        image_count = 0
        for item in structure:
            if item["type"] == "figure":
                image_count += 1
        return image_count
    
    def get_table_count(self, structure):
        """
        描述：获取表格数量
        参数：
            structure: 版面分析结果 List[Dict]
        返回值：
            table_count: 表格数量
        """
        table_count = 0
        for item in structure:
            if item["type"] == "table":
                table_count += 1
        return table_count

    def get_image_table_count_old(self, structure):
        '''
        描述：
            获取每一个文字块下方的图片数量和表格数量
        参数：  
            structure: 版面分析结果 List[Dict]
        返回值：    
            new_structure: 新的版面结构 List[Dict]
            Dict结构： {
                "content": 文字块内容,
                "image_count": 文字块下方的图片数量,
                "table_count": 文字块下方的表格数量
            }
        '''
        # 新的版面结构，包含每一块的类型、y坐标、[内容]
        new_structure = []
        for item in structure:
            axis_x_top  = item["bbox"][0]
            axis_y_top  = item["bbox"][1]
            axis_x_bottom = item["bbox"][2]
            axis_y_bottom = item["bbox"][3]
            axis_y_mean = np.mean([axis_y_top, axis_y_bottom])
            axis_x_mean = np.mean([axis_x_top, axis_x_bottom])
            tmp = {"type": item["type"], "axis_y_top": axis_y_top, "axis_x_top": axis_x_top, 
                    "axis_y_bottom": axis_y_bottom, "axis_x_bottom": axis_x_bottom, "axis_y_mean": axis_y_mean, "axis_x_mean": axis_x_mean}
            if item["type"] == "text":
                content = "".join([line["text"] for line in item["res"]])
                tmp["content"] = content
                new_structure.append(tmp)
            elif item["type"]  in ["table", "figure"]:
                new_structure.append(tmp)
        
        # 遍历每一块，获取阈值内的图片数量和表格数量
        for i, item in enumerate(new_structure):
            if item["type"] == "text":
                item["image_count"] = 0
                item["table_count"] = 0
                for j in range(i+1, len(new_structure)):
                    all_in = False
                    left_axis = 0
                    right_axis = 0
                    if abs(new_structure[j]["axis_x_top"]-item["axis_x_top"]) > 5:
                        left_axis = new_structure[j]["axis_x_top"]-item["axis_x_top"]
                    if abs(new_structure[j]["axis_x_bottom"]-item["axis_x_bottom"]) > 5:
                        right_axis = new_structure[j]["axis_x_bottom"]-item["axis_x_bottom"]
                    if left_axis * right_axis <= 0:
                        all_in = True    
                    # FIXME 下一块在当前块的上方，这个判断值为负数，需要修改
                    # FIXME 双列情况
                    if abs(new_structure[j]["axis_y_top"] - item["axis_y_bottom"]) < self.y_threshold and new_structure[j]["axis_x_mean"] - item["axis_x_mean"] < self.x_threshold:
                        # 如果下一块的顶部y坐标与当前块的底部y坐标的差小于阈值，且下一块的x轴与当前块的x轴相差小于阈值（在同一排）则计数
                        if new_structure[j]["type"] == "figure":
                            item["image_count"] += 1
                        elif new_structure[j]["type"] == "table":
                            item["table_count"] += 1
                    elif all_in:
                        if new_structure[j]["type"] == "figure":
                            item["image_count"] += 1
                        elif new_structure[j]["type"] == "table":
                            item["table_count"] += 1
                        # 如果下一块的顶部y坐标与当前块的底部y坐标的差大于阈值，则跳出循环
                    else:
                        break
        
        # 保留text类型的 content、image_count、table_count字段
        res_structure = [{k: item[k] for k in ["content", "image_count", "table_count"]} for item in new_structure if item["type"] == "text"]

        return res_structure

    def get_image_table_count(self, structure):
        '''
        描述：
            获取每一个文字块下方的图片数量和表格数量
            图表计算方式修改（段落<=2）
        参数：  
            structure: 版面分析结果 List[Dict]
        返回值：    
            new_structure: 新的版面结构 List[Dict]
            Dict结构： {
                "content": 文字块内容,
                "image_count": 文字块下方的图片数量,
                "table_count": 文字块下方的表格数量
            }
        '''
        # 新的版面结构，包含每一块的类型、y坐标、[内容]
        N = 1 # 遍历当前块下方的N块,如果不足N块，就到末尾
        new_structure = []
        for item in structure:
            axis_x_top  = item["bbox"][0]
            axis_y_top  = item["bbox"][1]
            axis_x_bottom = item["bbox"][2]
            axis_y_bottom = item["bbox"][3]
            axis_y_mean = np.mean([axis_y_top, axis_y_bottom])
            axis_x_mean = np.mean([axis_x_top, axis_x_bottom])
            tmp = {"type": item["type"], "axis_y_top": axis_y_top, "axis_x_top": axis_x_top, 
                    "axis_y_bottom": axis_y_bottom, "axis_x_bottom": axis_x_bottom, "axis_y_mean": axis_y_mean, "axis_x_mean": axis_x_mean}
            if item["type"] == "text":
                content = "".join([line["text"] for line in item["res"]])
                tmp["content"] = content
                new_structure.append(tmp)
            elif item["type"]  in ["table", "figure"]:
                new_structure.append(tmp)
        
        # 遍历每一块，获取阈值内的图片数量和表格数量
        for i, item in enumerate(new_structure):
            if item["type"] == "text":
                item["image_count"] = 0
                item["table_count"] = 0

                length = len(new_structure)
                # 如果当前块是最后一块，退出
                if i == length-1:
                    break

                # 遍历当前块下方的N块,如果不足N块，就到末尾
                for j in range(i+1, min(length, i+N+1)):
                    if new_structure[j]["type"] == "figure":
                        item["image_count"] += 1
                    elif new_structure[j]["type"] == "table":
                        item["table_count"] += 1
                
        # 保留text类型的 content、image_count、table_count字段
        res_structure = [{k: item[k] for k in ["content", "image_count", "table_count"]} for item in new_structure if item["type"] == "text"]

        return res_structure

    def get_text_bboxes(self, structure):
        """
        描述：获取文字的Bbox
        参数：
            content: 文字提取结果 List
        返回值：    
            bboxs: 文字的Bbox List
            Bbox: [左上角x，左上角y，右下角x，右下角y]
        """
        bboxs = []
        for line in structure:
            if line["type"] == "text":
                bboxs.append(line["bbox"])
        return bboxs

    def is_in_bboxes(self, bbox, bboxes, error_axis_x=50, error_axis_y=2):
        """
        描述：判断bbox是否在bboxes中
        参数：
            bbox: Bbox [左上角x，左上角y，右下角x，右下角y]
            bboxes: Bbox List
            error_axis_x: x轴允许误差
            error_axis_y: y轴允许误差
        返回值：
            True: 在bboxes中
            False: 不在bboxes中
        """
        # 添加允许误差
        bbox = [bbox[0]+error_axis_x, bbox[1]+error_axis_y, bbox[2]-error_axis_x, bbox[3]-error_axis_y]
        for box in bboxes:
            if bbox[0] >= box[0] - error_axis_x and bbox[1] >= box[1] - error_axis_y and bbox[2] <= box[2] + error_axis_x and bbox[3] <= box[3] + error_axis_y:
                return True
        return False

    def delete_image(self, imagepath):
        """
        描述:
            删除单张图片
        参数:
            imagepath:图片路径
        """
        os.remove(imagepath)

    def delete_images(self, filepaths):
        '''
        描述：
            批量删除图片
        参数：
            filepaths: List 文件路径
        '''
        for file in filepaths:
            os.remove(file)
    
    def is_single_colum(self, structure):
        """
        描述:
            判断是否为单栏页面
        参数:
            structure:该页面的结构化信息 List[Dict]
        返回值:
            true:单栏页面
            false:多栏页面
        """
        axis = 50 # 误差值
        middle_point = []
        for item in structure:
            middle_point.append(item["middle_point"])
        sorted_middle_point = sorted(middle_point, key= lambda t:t[1]) # 按照纵坐标升序
        for index in range(0,len(sorted_middle_point) - 1):
            temp1 = sorted_middle_point[index + 1][0] - sorted_middle_point[index][0] # 横坐标差值
            temp2 = sorted_middle_point[index + 1][1] - sorted_middle_point[index][1] # 纵坐标差值
            if temp2 <= axis and temp1 >= axis: # 横坐标差值大于误差且纵坐标差值小于误差，则判断为多栏
                return False
        return True

def clean_content(content):
    """
    描述：
        从文本中去除换行符、回车符、制表符、章节号
    参数：
        content: 文本内容
    返回值：    
        content: 处理后的文本内容
    """
    # 去除换行符、回车符、制表符
    content = content.replace('\n', '').replace('\r', '')
    content = content.replace('\t', '').replace(' ', '')
    # 去除所有章节号，例如4.3.1
    content = re.sub(r"\d+\.?\d*\.\d+\.?\d*\.\d+\.?\d*","", content)
    return content

def write_to_log(log_info, log_filepath='/code/scripts/log_mutil.txt'):
    """
    描述：
        将信息写入日志文件
    """
    now_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    with open(log_filepath, 'a', encoding='utf-8') as f:
        f.write(f"[{now_time}]: {log_info}")
